import argparse
import gzip
from deepsignal3.utils.process_utils import complement_seq
import numpy as np
import sys

def reverse_key(key_input):
    reverse_indexes=set()
    with open(key_input, 'r') as input_file:
        for line in input_file:
            key = line.strip()
            reverse_indexes.add(key)
    return reverse_indexes

def parse_args():
    parser = argparse.ArgumentParser("")
    parser.add_argument("--reverse", type=str,required=False)
    parser.add_argument("--input",'-i', type=str,required=True)
    parser.add_argument("--output",'-o', type=str,required=True)
    return parser.parse_args()

def main():
    args=parse_args()
    if args.input.endswith(".gz"):
        input_file = gzip.open(args.input, 'rt')
    else:
        input_file = open(args.input, 'r')

    if args.reverse is not None:
        reverse_indexes=reverse_key(args.reverse)

    with open(args.output, 'w') as output_file:
        for line in input_file:
            words=line.strip().split('\t')
            index=words[0]
            if index in reverse_indexes:
                seq_len=len(words[3])
                words[3]=complement_seq(words[3])
                pred_pos=np.array([int(x) for x in words[5].split(",")])
                for i in range(pred_pos.shape[0]):
                    pos=seq_len-int(pred_pos[i])-2
                    if words[3][pos]!='C':
                        KeyError('convert error!!!')
                    pred_pos[i]=pos
                words[5]=','.join([str(x) for x in pred_pos])
            processed_line = '\t'.join(words)
            output_file.write(processed_line+'\n')
if __name__ == '__main__':
    sys.exit(main())
            